import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
import logging
from tqdm import tqdm


class DACDataset(Dataset):
    """针对Criteo Display Advertising Challenge (dac-v1)数据集的自定义数据集类"""

    # self.data 原始数据数据
    # self.num_cols_indices 连续类型特征索引
    # self.cat_cols_indices 离散型特征索引
    # self.target_col_index 目标值索引
    # num_scaler 数据处理 可自定义 MinMaxScaler()用于将数据转换为0~1之间的数据
    # self.cat_encoders 离散数据枚举 le进行列表编码后的结果
    # self.cat_dims 每个离散纬度的枚举个数

    def __init__(self, file_path, sample_size=None):
        self.file_path = file_path
        self.sample_size = sample_size

        # 读取数据
        logging.info(f"开始读取数据文件: {file_path}")
        self.data = pd.read_csv(
            file_path,
            sep='\t',
            header=None,
            nrows=sample_size,
            low_memory=False
        )
        logging.info(f"数据读取完成，形状: {self.data.shape}")

        # 验证列数
        expected_cols = 40
        if self.data.shape[1] != expected_cols:
            raise ValueError(f"数据列数不匹配: 预期{expected_cols}列，实际{self.data.shape[1]}列")

        # 填充缺失值
        self.data = self.data.fillna({i: 'NaN' for i in range(self.data.shape[1])})

        # 特征索引配置
        self.num_cols_indices = list(range(1, 14))  # 1-13列：整数特征
        self.cat_cols_indices = list(range(14, expected_cols))  # 14列及以后：类别特征
        self.target_col_index = 0  # 训练集第0列是目标值

        # 处理数值特征
        if self.num_cols_indices:
            self.num_scaler = MinMaxScaler()
            # 替换'NaN'为0并转换为float
            self.data[self.num_cols_indices] = self.data[self.num_cols_indices].replace('NaN', 0).infer_objects(
                copy=False)
            self.data[self.num_cols_indices] = self.data[self.num_cols_indices].astype(float)
            self.data[self.num_cols_indices] = self.num_scaler.fit_transform(self.data[self.num_cols_indices])
        else:
            self.num_scaler = None

        self.cat_encoders = {}
        for idx in tqdm(self.cat_cols_indices, desc="拟合类别特征编码器"):
            le = LabelEncoder()
            self.data[idx] = le.fit_transform(self.data[idx].astype(str))
            self.cat_encoders[idx] = le

        # 类别特征维度
        self.cat_dims = [len(le.classes_) for le in self.cat_encoders.values()]
        logging.info(
            f"特征处理完成 - 数值特征: {len(self.num_cols_indices)}个, 类别特征: {len(self.cat_cols_indices)}个")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = self.data.iloc[idx]

        # 提取数值特征
        num_features = torch.tensor([row[i] for i in self.num_cols_indices], dtype=torch.float32)

        # 提取类别特征
        cat_features = torch.tensor([row[i] for i in self.cat_cols_indices], dtype=torch.long)
        target = torch.tensor(row[self.target_col_index], dtype=torch.float32)
        return cat_features, num_features, target
